{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "h2q27gKz1H20"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "TUfAcER1oUS6"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Gb7qyhNL1yWt"
      },
      "source": [
        "# Image classification with TensorFlow Lite Model Maker with TensorFlow 2.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "nDABAblytltI"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/examples/blob/master/tensorflow_examples/lite/model_maker/demo/image_classification.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003e\n",
        "    Run in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/examples/blob/master/tensorflow_examples/lite/model_maker/demo/image_classification.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003e\n",
        "    View source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "m86-Nh4pMHqY"
      },
      "source": [
        "Model Maker library simplifies the process of adapting and converting a TensorFlow neural-network model to particular input data when deploying this model for on-device ML applications.\n",
        "\n",
        "This notebook shows an end-to-end example that utilizes this Model Maker library to illustrate the adaption and conversion of a commonly-used image classification model to classify flowers on a mobile device."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bcLF2PKkSbV3"
      },
      "source": [
        "## Prerequisites\n",
        "\n",
        "To run this example, we first need to install serveral required packages, including Model Maker package that in github [repo](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6cv3K3oaksJv"
      },
      "outputs": [],
      "source": [
        "!pip install git+git://github.com/tensorflow/examples.git#egg=tensorflow-examples[model_maker,metadata]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Gx1HGRoFQ54j"
      },
      "source": [
        "Import the required packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "XtxiUeZEiXpt"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "import tensorflow as tf\n",
        "assert tf.__version__.startswith('2')\n",
        "\n",
        "from tensorflow_examples.lite.model_maker.core.data_util.image_dataloader import ImageClassifierDataLoader\n",
        "from tensorflow_examples.lite.model_maker.core.task import image_classifier\n",
        "from tensorflow_examples.lite.model_maker.core.task.model_spec import mobilenet_v2_spec\n",
        "from tensorflow_examples.lite.model_maker.core.task.model_spec import ImageModelSpec\n",
        "\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "KKRaYHABpob5"
      },
      "source": [
        "## Simple End-to-End Example"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SiZZ5DHXotaW"
      },
      "source": [
        "### Get the data path\n",
        "\n",
        "Let's get some images to play with this simple end-to-end example. Hundreds of images is a good start for Model Maker while more data could achieve better accuracy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "3jz5x0JoskPv"
      },
      "outputs": [],
      "source": [
        "image_path = tf.keras.utils.get_file(\n",
        "      'flower_photos',\n",
        "      'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',\n",
        "      untar=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "a55MR6i6nuDm"
      },
      "source": [
        "You could replace `image_path` with your own image folders. As for uploading data to colab, you could find the upload button in the left sidebar shown in the image below with the red rectangle. Just have a try to upload a zip file and unzip it. The root file path is the current path.\n",
        "\n",
        "\u003cimg src=\"https://storage.googleapis.com/download.tensorflow.org/models/tflite/screenshots/model_maker_image_classification.png\" alt=\"Upload File\" width=\"800\" hspace=\"100\"\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NNRNv_mloS89"
      },
      "source": [
        "If you prefer not to upload your images to the cloud, you could try to run the library locally following the [guide](https://github.com/tensorflow/examples/tree/master/tensorflow_examples/lite/model_maker) in github."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "w-VDriAdsowu"
      },
      "source": [
        "### Run the example\n",
        "The example just consists of 4 lines of code as shown below, each of which representing one step of the overall process.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6ahtcO86tZBL"
      },
      "source": [
        "1.   Load input data specific to an on-device ML app. Split it to training data and testing data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "lANoNS_gtdH1"
      },
      "outputs": [],
      "source": [
        "data = ImageClassifierDataLoader.from_folder(image_path)\n",
        "train_data, test_data = data.split(0.9)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Y_9IWyIztuRF"
      },
      "source": [
        "2. Customize the TensorFlow model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yRXMZbrwtyRD"
      },
      "outputs": [],
      "source": [
        "model = image_classifier.create(train_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "oxU2fDr-t2Ya"
      },
      "source": [
        "3. Evaluate the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "wQr02VxJt6Cs"
      },
      "outputs": [],
      "source": [
        "loss, accuracy = model.evaluate(test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "eVZw9zU8t84y"
      },
      "source": [
        "4.  Export to TensorFlow Lite model.\n",
        "\n",
        "Here, we export TensorFlow Lite model with [metadata](https://www.tensorflow.org/lite/convert/metadata) which provides a standard for model descriptions.\n",
        "You could download it in the left sidebar same as the uploading part for your own use."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Zb-eIzfluCoa"
      },
      "outputs": [],
      "source": [
        "model.export('image_classifier.tflite', 'image_labels.txt', with_metadata=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pyju1qc_v-wy"
      },
      "source": [
        "After this simple 4 steps, we could further use TensorFlow Lite model file and label file in on-device applications like in [image classification](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification) reference app.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "R1QG32ivs9lF"
      },
      "source": [
        "## Detailed Process\n",
        "\n",
        "Currently, we support several models such as  EfficientNet-Lite* models, MobileNetV2, ResNet50 as pre-trained models for image classification. But it is very flexible to add new pre-trained models to this library with just a few lines of code.\n",
        "\n",
        "\n",
        "The following walks through this end-to-end example step by step to show more detail."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ygEncJxtl-nQ"
      },
      "source": [
        "### Step 1: Load Input Data Specific to an On-device ML App\n",
        "\n",
        "The flower dataset contains 3670 images belonging to 5 classes. Download the archive version of the dataset and untar it.\n",
        "\n",
        "The dataset has the following directory structure:\n",
        "\n",
        "\u003cpre\u003e\n",
        "\u003cb\u003eflower_photos\u003c/b\u003e\n",
        "|__ \u003cb\u003edaisy\u003c/b\u003e\n",
        "    |______ 100080576_f52e8ee070_n.jpg\n",
        "    |______ 14167534527_781ceb1b7a_n.jpg\n",
        "    |______ ...\n",
        "|__ \u003cb\u003edandelion\u003c/b\u003e\n",
        "    |______ 10043234166_e6dd915111_n.jpg\n",
        "    |______ 1426682852_e62169221f_m.jpg\n",
        "    |______ ...\n",
        "|__ \u003cb\u003eroses\u003c/b\u003e\n",
        "    |______ 102501987_3cdb8e5394_n.jpg\n",
        "    |______ 14982802401_a3dfb22afb.jpg\n",
        "    |______ ...\n",
        "|__ \u003cb\u003esunflowers\u003c/b\u003e\n",
        "    |______ 12471791574_bb1be83df4.jpg\n",
        "    |______ 15122112402_cafa41934f.jpg\n",
        "    |______ ...\n",
        "|__ \u003cb\u003etulips\u003c/b\u003e\n",
        "    |______ 13976522214_ccec508fe7.jpg\n",
        "    |______ 14487943607_651e8062a1_m.jpg\n",
        "    |______ ...\n",
        "\u003c/pre\u003e"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "7tOfUr2KlgpU"
      },
      "outputs": [],
      "source": [
        "image_path = tf.keras.utils.get_file(\n",
        "      'flower_photos',\n",
        "      'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',\n",
        "      untar=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "E051HBUM5owi"
      },
      "source": [
        "Use `ImageClassifierDataLoader` class to load data.\n",
        "\n",
        "As for `from_folder()` method, it could load data from the folder. It assumes that the image data of the same class are in the same subdirectory and the subfolder name is the class name. Currently, JPEG-encoded images and PNG-encoded images are supported."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "I_fOlZsklmlL"
      },
      "outputs": [],
      "source": [
        "data = ImageClassifierDataLoader.from_folder(image_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "u501eT4koURB"
      },
      "source": [
        "Split it to training data (80%), validation data (10%, optional) and testing data (10%)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cY4UU5SUobtJ"
      },
      "outputs": [],
      "source": [
        "train_data, rest_data = data.split(0.8)\n",
        "validation_data, test_data = rest_data.split(0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Z9_MYPie3EMO"
      },
      "source": [
        "Show 25 image examples with labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Ih4Wx44I482b"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(10,10))\n",
        "for i, (image, label) in enumerate(data.dataset.take(25)):\n",
        "  plt.subplot(5,5,i+1)\n",
        "  plt.xticks([])\n",
        "  plt.yticks([])\n",
        "  plt.grid(False)\n",
        "  plt.imshow(image.numpy(), cmap=plt.cm.gray)\n",
        "  plt.xlabel(data.index_to_label[label.numpy()])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "AWuoensX4vDA"
      },
      "source": [
        "### Step 2: Customize the TensorFlow Model\n",
        "\n",
        "Create a custom image classifier model based on the loaded data. The default model is EfficientNet-Lite0.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "TvYSUuJY3QxR"
      },
      "outputs": [],
      "source": [
        "model = image_classifier.create(train_data, validation_data=validation_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4JFOKWnH9x8_"
      },
      "source": [
        "Have a look at the detailed model structure."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "QNXAfjl192dC"
      },
      "outputs": [],
      "source": [
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LP5FPk_tOxoZ"
      },
      "source": [
        "### Step 3: Evaluate the Customized Model\n",
        "\n",
        "Evaluate the result of the model, get the loss and accuracy of the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "A8c2ZQ0J3Riy"
      },
      "outputs": [],
      "source": [
        "loss, accuracy = model.evaluate(test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6ZCrYOWoCt05"
      },
      "source": [
        "We could plot the predicted results in 100 test images. Predicted labels with red color are the wrong predicted results while others are correct."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "n9O9Kx7nDQWD"
      },
      "outputs": [],
      "source": [
        "# A helper function that returns 'red'/'black' depending on if its two input\n",
        "# parameter matches or not.\n",
        "def get_label_color(val1, val2):\n",
        "  if val1 == val2:\n",
        "    return 'black'\n",
        "  else:\n",
        "    return 'red'\n",
        "\n",
        "# Then plot 100 test images and their predicted labels.\n",
        "# If a prediction result is different from the label provided label in \"test\"\n",
        "# dataset, we will highlight it in red color.\n",
        "plt.figure(figsize=(20, 20))\n",
        "predicts = model.predict_top_k(test_data)\n",
        "for i, (image, label) in enumerate(test_data.dataset.take(100)):\n",
        "  ax = plt.subplot(10, 10, i+1)\n",
        "  plt.xticks([])\n",
        "  plt.yticks([])\n",
        "  plt.grid(False)\n",
        "  plt.imshow(image.numpy(), cmap=plt.cm.gray)\n",
        "\n",
        "  predict_label = predicts[i][0][0]\n",
        "  color = get_label_color(predict_label,\n",
        "                          test_data.index_to_label[label.numpy()])\n",
        "  ax.xaxis.label.set_color(color)\n",
        "  plt.xlabel('Predicted: %s' % predict_label)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "S3H0rkbLUZAG"
      },
      "source": [
        "If the accuracy doesn't meet the app requirement, one could refer to [Advanced Usage](#scrollTo=zNDBP2qA54aK) to explore alternatives such as changing to a larger model, adjusting re-training parameters etc."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aeHoGAceO2xV"
      },
      "source": [
        "### Step 4: Export to TensorFlow Lite Model\n",
        "\n",
        "Convert the existing model to TensorFlow Lite model format and save the image labels in label file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Im6wA9lK3TQB"
      },
      "outputs": [],
      "source": [
        "model.export('flower_classifier.tflite', 'flower_labels.txt', with_metadata=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ROS2Ay2jMPCl"
      },
      "source": [
        "The TensorFlow Lite model file and label file could be used in [image classification](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification) reference app.\n",
        "\n",
        "As for android reference app as an example, we could add `flower_classifier.tflite` and `flower_label.txt` in [assets](https://github.com/tensorflow/examples/tree/master/lite/examples/image_classification/android/app/src/main/assets) folder. Meanwhile, change label filename in [code](https://github.com/tensorflow/examples/blob/master/lite/examples/image_classification/android/app/src/main/java/org/tensorflow/lite/examples/classification/tflite/ClassifierFloatMobileNet.java#L65) and TensorFlow Lite file name in [code](https://github.com/tensorflow/examples/blob/master/lite/examples/image_classification/android/app/src/main/java/org/tensorflow/lite/examples/classification/tflite/ClassifierFloatMobileNet.java#L60). Thus, we could run the retrained float TensorFlow Lite model on the android app.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-4jQaxyT5_KV"
      },
      "source": [
        "Here, we also demonstrate how to use the above files to run and evaluate the TensorFlow Lite model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "S1YoPX5wOK-u"
      },
      "outputs": [],
      "source": [
        "# Read TensorFlow Lite model from TensorFlow Lite file.\n",
        "with tf.io.gfile.GFile('flower_classifier.tflite', 'rb') as f:\n",
        "  model_content = f.read()\n",
        "\n",
        "# Read label names from label file.\n",
        "with tf.io.gfile.GFile('flower_labels.txt', 'r') as f:\n",
        "  label_names = f.read().split('\\n')\n",
        "\n",
        "# Initialze TensorFlow Lite inpterpreter.\n",
        "interpreter = tf.lite.Interpreter(model_content=model_content)\n",
        "interpreter.allocate_tensors()\n",
        "input_index = interpreter.get_input_details()[0]['index']\n",
        "output = interpreter.tensor(interpreter.get_output_details()[0][\"index\"])\n",
        "\n",
        "# Run predictions on each test image data and calculate accuracy.\n",
        "accurate_count = 0\n",
        "for i, (image, label) in enumerate(test_data.dataset):\n",
        "    # Pre-processing should remain the same. Currently, just normalize each pixel value and resize image according to the model's specification.\n",
        "    image, _ = model.preprocess(image, label)\n",
        "    # Add batch dimension and convert to float32 to match with the model's input\n",
        "    # data format.\n",
        "    image = tf.expand_dims(image, 0).numpy()\n",
        "\n",
        "    # Run inference.\n",
        "    interpreter.set_tensor(input_index, image)\n",
        "    interpreter.invoke()\n",
        "\n",
        "    # Post-processing: remove batch dimension and find the label with highest\n",
        "    # probability.\n",
        "    predict_label = np.argmax(output()[0])\n",
        "    # Get label name with label index.\n",
        "    predict_label_name = label_names[predict_label]\n",
        "\n",
        "    accurate_count += (predict_label == label.numpy())\n",
        "\n",
        "accuracy = accurate_count * 1.0 / test_data.size\n",
        "print('TensorFlow Lite model accuracy = %.4f' % accuracy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "fuHB-NFqpKTD"
      },
      "source": [
        "Note that preprocessing for inference should be the same as training. Currently, preprocessing contains normalizing each pixel value and resizing the image according to the model's specification. For  EfficientNet-Lite0, input image should be normalized to `[0, 1]` and resized to `[224, 224, 3]`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zNDBP2qA54aK"
      },
      "source": [
        "## Advanced Usage\n",
        "\n",
        "The `create` function is the critical part of this library. It uses transfer learning with a pretrained model similiar to the [tutorial](https://www.tensorflow.org/tutorials/images/transfer_learning).\n",
        "\n",
        "The `create`function contains the following steps:\n",
        "\n",
        "1.   Split the data into training, validation, testing data according to parameter `validation_ratio` and `test_ratio`. The default value of `validation_ratio` and `test_ratio` are `0.1` and `0.1`.\n",
        "2.   Download a [Image Feature Vector](https://www.tensorflow.org/hub/common_signatures/images#image_feature_vector) as the base model from TensorFlow Hub. The default pre-trained model is  EfficientNet-Lite0.\n",
        "3.   Add a classifier head with a Dropout Layer with `dropout_rate` between head layer and pre-trained model. The default `dropout_rate` is the default `dropout_rate` value from [make_image_classifier_lib](https://github.com/tensorflow/hub/blob/master/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py#L55) by TensorFlow Hub.\n",
        "4.   Preprocess the raw input data. Currently, preprocessing steps including normalizing the value of each image pixel to model input scale and resizing it to model input size.   EfficientNet-Lite0 have the input scale `[0, 1]` and the input image size `[224, 224, 3]`.\n",
        "5.   Feed the data into the classifier model. By default, the training parameters such as training epochs, batch size, learning rate, momentum are the default values from [make_image_classifier_lib](https://github.com/tensorflow/hub/blob/master/tensorflow_hub/tools/make_image_classifier/make_image_classifier_lib.py#L55) by TensorFlow Hub. Only the classifier head is trained.\n",
        "\n",
        "\n",
        "In this section, we describe several advanced topics, including switching to a different image classification model, changing the training hyperparameters etc.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "A4kiTJtZ_sDm"
      },
      "source": [
        "## Change the model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "794vgj6ud7Ep"
      },
      "source": [
        "### Change to the model that's supported in this library.\n",
        "\n",
        "This library supports  EfficientNet-Lite models, MobileNetV2, ResNet50 by now. [EfficientNet-Lite](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/lite) are a family of image classification models that could achieve state-of-art accuracy and suitable for Edge devices. The default model is EfficientNet-Lite0.\n",
        "\n",
        "We could switch model to MobileNetV2 by just setting parameter `model_spec` to  `mobilenet_v2_spec` in `create` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "7JKsJ6-P6ae1"
      },
      "outputs": [],
      "source": [
        "model = image_classifier.create(train_data, model_spec=mobilenet_v2_spec, validation_data=validation_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gm_B1Wv08AxR"
      },
      "source": [
        "Evaluate the newly retrained MobileNetV2 model to see the accuracy and loss in testing data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "lB2Go3HW8X7_"
      },
      "outputs": [],
      "source": [
        "loss, accuracy = model.evaluate(test_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vAciGzVWtmWp"
      },
      "source": [
        "### Change to the model in TensorFlow Hub\n",
        "\n",
        "Moreover, we could also switch to other new models that inputs an image and outputs a feature vector with TensorFlow Hub format.\n",
        "\n",
        "As [Inception V3](https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1) model as an example, we could define `inception_v3_spec` which is an object of `ImageModelSpec` and contains the specification of the Inception V3 model.\n",
        "\n",
        "We need to specify the model name `name`, the url of the TensorFlow Hub model `uri`. Meanwhile, the default value of `input_image_shape` is `[224, 224]`. We need to change it to `[299, 299]` for Inception V3 model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "xdiMF2WMfAR4"
      },
      "outputs": [],
      "source": [
        "inception_v3_spec = ImageModelSpec(\n",
        "    uri='https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1')\n",
        "inception_v3_spec.input_image_shape = [299, 299]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "T_GGIoXZCs5F"
      },
      "source": [
        "Then, by setting parameter `model_spec` to `inception_v3_spec` in `create` method, we could retrain the Inception V3 model. \n",
        "\n",
        "The remaining steps are exactly same and we could get a customized InceptionV3 TensorFlow Lite model in the end."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UhZ5IRKdeex3"
      },
      "source": [
        "### Change your own custom model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "svTjlZhrCrcV"
      },
      "source": [
        "If we'd like to use the custom model that's not in TensorFlow Hub, we should create and export [ModelSpec](https://www.tensorflow.org/hub/api_docs/python/hub/ModuleSpec) in TensorFlow Hub.\n",
        "\n",
        "Then start to define `ImageModelSpec` object like the process above."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4M9bn703AHt2"
      },
      "source": [
        "## Change the training hyperparameters\n",
        "We could also change the training hyperparameters like `epochs`, `dropout_rate` and `batch_size` that could affect the model accuracy. For instance,\n",
        "\n",
        "\n",
        "*   `epochs`: more epochs could achieve better accuracy until it converges but training for too many epochs may lead to overfitting.\n",
        "*   `dropout_rate`: avoid overfitting.\n",
        "*   `batch_size`: number of samples to use in one training step.\n",
        "*   `validation_data`: number of samples to use in one training step.\n",
        "\n",
        "\n",
        "For example, we could train with more epochs.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "A3k7mhH54QcK"
      },
      "outputs": [],
      "source": [
        "model = image_classifier.create(train_data, validation_data=validation_data, epochs=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VaYBQymQDsXU"
      },
      "source": [
        "Evaluate the newly retrained model with 10 training epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VafIYpKWD4Sw"
      },
      "outputs": [],
      "source": [
        "loss, accuracy = model.evaluate(test_data)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "image_classification.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
